#!/usr/bin/env python
#-*- coding: utf-8 -*-

import re, sys


def fits(x, n):
  return -(2**(n-1)) <= x <= (2**(n-1)-1)


def ld_w(m):
  regA = m.group(1).strip()
  regB = m.group(2).strip()

  try:
    val  = int(m.group(3).strip())
    if regA == regB or regB.lower() == '$zero' or fits(val, 11):
      return 'ld %s, [%s + %d]' % (regA, regB, val)
  except:
      pass

  return 'mov(%%tmp, %s); ld %s, [%s + %%tmp]' % (val, regA, regB)


def st_w(m):
  regA = m.group(1).strip()
  regB = m.group(2).strip()

  try:
    val  = int(m.group(3).strip())
    if regA == regB or regB.lower() == '$zero' or fits(val, 11):
      return 'st [%s + %d], %s' % (regB, val, regA)
  except:
    pass

  return 'mov (%%tmp, %s); st [%s + %%tmp], %s' % (val, regB, regA)


def slti(m):
  regA = m.group(1).strip()
  regB = m.group(2).strip()
  val  = int(m.group(3).strip())
  
  hd = ''
  
  if regA == regB or regB.lower() == '$zero' or fits(val, 11):
    hd = 'cmp (%s, %d); ' % (regB, val)
  else:
    hd = 'mov (%%tmp, %d); cmp (%s, %%tmp); ' % (val, regB)
  
  return hd + 'and %s, $psw, LT_MASK; sra %s, %s, LT_POS \t# SLTI' % (regA, regA, regA)
         

def sltiu(m):
  regA = m.group(1).strip()
  regB = m.group(2).strip()
  val  = int(m.group(3).strip())
  
  hd = ''
  
  if regB == "$zero" or fits(val, 11):
    hd = 'sub! $zero, %s, %d; ' % (regB, val)
  else:
    hd = 'mov (%%tmp, %d); sub! $zero, %s, %%tmp; ' % (val, regB)
  
  return hd + 'and %s, $psw, N_MASK; srl %s, %s, N_POS \t# SLTIU' % (regA, regA, regA)


def sethi(m):
  dst = m.group(1).strip()
  imm = m.group(3).strip()

  return 'add %s, %s' % (dst, imm)


def repl_ri(op):
  def func(m):
    regA = m.group(1).strip()
    regB = m.group(2).strip()

    val  = m.group(3).strip()

    try:
      val = int(val)
      if regA == regB or regB.lower() == '$zero' or fits(val, 11):
        return '%s %s, %s, %s' % (op, regA, regB, val)
    except:
      pass

    return 'mov (%%tmp, %s); %s %s, %s, %%tmp' % (val, op, regA, regB)
      
  return func
  

if __name__ == '__main__':
  directives = ''
  rest = ''
  
  with open(sys.argv[1], 'r') as f:
   headers = ['.global', '.extern', '#include']
   exclude = ['# .align', '.weak']

   for line in f.readlines():
     line_s = line.strip()
     
     if any(line_s.startswith(hd) for hd in headers):
       directives += line
     elif any(line_s.startswith(hd) for hd in exclude):
       pass
     elif re.match(r'^((\.section(\s+)\.rodata)|(\.bss))', line_s):
       rest += '\t.data\n' 
     elif re.match(r'^.section\s+(\.ctors|\.gnu)', line_s):
       rest += '\t.text\n'
     else:
       rest += line
    
  source = directives + rest
  source = re.sub(r'@ld_w (.*), (.*), ([^;#\s]*)', ld_w, source)
  source = re.sub(r'@st_w (.*), (.*), ([^;#\s]*)', st_w, source)
  
  source = re.sub(r'@add (.*), (.*), ([^;#\s]*)', repl_ri('add'), source)
  source = re.sub(r'@and (.*), (.*), ([^;#\s]*)', repl_ri('and'), source)
  source = re.sub(r'@or (.*), (.*), ([^;#\s]*)', repl_ri('or'), source)
  source = re.sub(r'@xor (.*), (.*), ([^;#\s]*)', repl_ri('xor'), source)
  
  source = re.sub(r'@slti (.*), (.*), ([^;#\s]*)', slti, source)
  source = re.sub(r'@sltiu (.*), (.*), ([^;#\s]*)', sltiu, source)

  with open(sys.argv[1], 'w') as f:
    f.write(source)
